// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <utility>



namespace nebula {
namespace internal {

/// Helper providing single-lookup conditional insertion into std::map or
/// std::unordered_map. If `key` exists in the container, an iterator to that pair
/// will be returned. If `key` does not exist in the container, `gen(key)` will be
/// invoked and its return value inserted.
template <typename Map, typename Gen>
auto GetOrInsertGenerated(Map* map, typename Map::key_type key, Gen&& gen)
    -> decltype(map->begin()->second = gen(map->begin()->first), map->begin()) {
  decltype(gen(map->begin()->first)) placeholder{};

  auto it_success = map->emplace(std::move(key), std::move(placeholder));
  if (it_success.second) {
    // insertion of placeholder succeeded, overwrite it with gen()
    const auto& inserted_key = it_success.first->first;
    auto* value = &it_success.first->second;
    *value = gen(inserted_key);
  }
  return it_success.first;
}

template <typename Map, typename Gen>
auto GetOrInsertGenerated(Map* map, typename Map::key_type key, Gen&& gen)
    -> turbo::Result<decltype(map->begin()->second = gen(map->begin()->first).value_or_die(),
                       map->begin())> {
  decltype(gen(map->begin()->first).value_or_die()) placeholder{};

  auto it_success = map->emplace(std::move(key), std::move(placeholder));
  if (it_success.second) {
    // insertion of placeholder succeeded, overwrite it with gen()
    const auto& inserted_key = it_success.first->first;
    auto* value = &it_success.first->second;
    TURBO_MOVE_OR_RAISE(*value, gen(inserted_key));
  }
  return it_success.first;
}

}  // namespace internal
}  // namespace nebula
